import os
from PIL import Image
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
from reportlab.lib.styles import ParagraphStyle
from reportlab.platypus import Paragraph
from reportlab.lib.utils import ImageReader  # 导入ImageReader
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
import base64
from openai import OpenAI
import io
import time
import sys

# 指定目录路径 - 使用脚本所在目录，避免权限问题
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))  # 获取脚本所在目录
IMAGE_DIRECTORY = SCRIPT_DIR  # 使用脚本所在目录作为图片目录
OUTPUT_PDF = os.path.join(SCRIPT_DIR, "图片描述结果.pdf")  # 在同一目录下保存PDF

# 豆包API配置 - 使用火山方舟API（修正域名）
API_KEY = os.environ.get("doubao_API_KEY")  # 您的API密钥
# 修正了域名拼写：volces -> volcengine
BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"  # 正确的API端点
MODEL_NAME = "doubao-1-5-vision-pro-32k-250115"  # 豆包模型名称

# 最大超时时间（秒）
REQUEST_TIMEOUT = 30
# 最大重试次数
MAX_RETRIES = 3
# 最大图片大小（宽度和高度，像素）
MAX_IMAGE_SIZE = 1024


def setup_chinese_font():
    """设置中文字体支持"""
    try:
        # 尝试加载不同的中文字体，提高兼容性
        font_options = [
            ('SimSun', 'STKAITI.ttf'),
            ('SimSun', 'simsun.ttc'),
            ('SimHei', 'simhei.ttf'),
            ('Microsoft YaHei', 'msyh.ttf')
        ]

        for font_name, font_file in font_options:
            try:
                # 首先尝试从当前目录加载
                if os.path.exists(os.path.join(SCRIPT_DIR, font_file)):
                    font_path = os.path.join(SCRIPT_DIR, font_file)
                # 然后尝试从系统字体目录加载（Windows）
                elif os.path.exists(os.path.join("C:\\Windows\\Fonts", font_file)):
                    font_path = os.path.join("C:\\Windows\\Fonts", font_file)
                else:
                    continue

                pdfmetrics.registerFont(TTFont(font_name, font_path))
                print(f"成功加载字体: {font_name}")
                return font_name
            except Exception as e:
                print(f"加载字体 {font_name} 失败: {str(e)}")
                continue

        print("警告: 无法加载任何中文字体，PDF中的中文可能无法正确显示")
        return None
    except Exception as e:
        print(f"设置中文字体时出错: {str(e)}")
        return None


def resize_image_if_needed(image_path, max_size=MAX_IMAGE_SIZE):
    """调整图片大小以优化API处理"""
    try:
        with Image.open(image_path) as img:
            # 检查图片是否需要调整大小
            if img.width > max_size or img.height > max_size:
                # 计算缩放比例
                ratio = min(max_size / img.width, max_size / img.height)
                new_width = int(img.width * ratio)
                new_height = int(img.height * ratio)

                # 调整图片大小
                resized_img = img.resize((new_width, new_height), Image.LANCZOS)

                # 创建内存中的图片对象
                img_byte_array = io.BytesIO()
                # 保存为JPEG格式，优化文件大小
                resized_img.save(img_byte_array, format='JPEG', quality=85)
                img_byte_array.seek(0)

                print(f"已将图片 {os.path.basename(image_path)} 调整为 {new_width}x{new_height}")
                return img_byte_array.getvalue()
            else:
                # 如果图片尺寸合适，仍然转换为JPEG以减小大小
                if img.format != 'JPEG':
                    img_byte_array = io.BytesIO()
                    img.save(img_byte_array, format='JPEG', quality=85)
                    img_byte_array.seek(0)
                    print(f"已将图片 {os.path.basename(image_path)} 转换为JPEG格式以优化大小")
                    return img_byte_array.getvalue()
                else:
                    with open(image_path, "rb") as f:
                        return f.read()
    except Exception as e:
        print(f"调整图片大小时出错 ({image_path}): {str(e)}")
        # 如果处理失败，返回原始图片
        with open(image_path, "rb") as f:
            return f.read()


def encode_image_to_base64(image_data):
    """将图片数据转换为Base64编码"""
    try:
        return base64.b64encode(image_data).decode('utf-8')
    except Exception as e:
        print(f"图片编码失败: {str(e)}")
        return None


def analyze_image_with_doubao(image_path):
    """使用豆包API分析图片内容"""
    if not API_KEY:
        return "错误: 未设置API密钥。请设置环境变量 'doubao_API_KEY'"

    # 重试机制
    for attempt in range(MAX_RETRIES):
        try:
            # 优化图片大小
            image_data = resize_image_if_needed(image_path)
            base64_image = encode_image_to_base64(image_data)

            if not base64_image:
                return f"无法处理图片 {os.path.basename(image_path)}: 图片编码失败"

            print(f"正在分析图片: {os.path.basename(image_path)} (尝试 {attempt + 1}/{MAX_RETRIES})")

            client = OpenAI(
                base_url=BASE_URL,
                api_key=API_KEY
            )

            # 添加超时设置
            response = client.chat.completions.create(
                model=MODEL_NAME,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": "请分析这张图片并按以下格式描述：\n1. 图中共有几个人，都在做什么\n2. 其中老人有几个\n3. 有老人摔倒吗\n4. 用500字左右写一篇看图作文，加点想象"
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpeg;base64,{base64_image}"
                                }
                            }
                        ]
                    }
                ],
                max_tokens=500,
                timeout=REQUEST_TIMEOUT  # 设置请求超时
            )

            description = response.choices[0].message.content
            print(f"图片分析成功: {os.path.basename(image_path)}")
            return description

        except KeyboardInterrupt:
            print("\n操作被用户中断")
            sys.exit(1)
        except Exception as e:
            print(f"调用API时发生错误 (尝试 {attempt + 1}/{MAX_RETRIES}): {str(e)}")
            if attempt < MAX_RETRIES - 1:
                wait_time = 2 ** attempt  # 指数退避
                print(f"等待 {wait_time} 秒后重试...")
                time.sleep(wait_time)
            else:
                return f"在 {MAX_RETRIES} 次尝试后分析图片失败: {str(e)}"

    return "分析图片失败，达到最大重试次数"


def create_pdf_with_images(image_descriptions, output_pdf, chinese_font=None):
    """生成包含图片和描述的PDF文件"""
    if not image_descriptions:
        print("没有找到有效的图片描述，不创建PDF")
        return

    try:
        c = canvas.Canvas(output_pdf, pagesize=letter)
        width, height = letter

        # 如果没有中文字体，使用默认字体
        font_name = chinese_font if chinese_font else "Helvetica"

        # 定义段落样式
        styles = {
            'default': ParagraphStyle(
                'default',
                fontName=font_name,
                fontSize=12,
                leading=14,  # 行间距
            )
        }

        # 添加封面
        c.setFont(font_name, 24)
        c.drawString(50, height - 100, "图片场景分析报告")
        c.setFont(font_name, 14)
        c.drawString(50, height - 150, f"总共分析图片: {len(image_descriptions)} 张")
        c.drawString(50, height - 180, f"生成时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
        c.showPage()

        # 添加每张图片及其描述
        for img_path, description in image_descriptions:
            try:
                c.showPage()

                y_position = height - 30
                c.setFont(font_name, 16)
                c.drawString(50, y_position, "图片场景分析")
                y_position -= 30

                img_name = os.path.basename(img_path)
                c.setFont(font_name, 14)
                c.drawString(50, y_position, f"图片: {img_name}")
                y_position -= 30

                try:
                    img = Image.open(img_path)
                    img_width, img_height = img.size
                    scale_ratio = min((width - 100) / img_width, 300 / img_height)
                    new_width = img_width * scale_ratio
                    new_height = img_height * scale_ratio
                    # 使用ImageReader读取图片
                    c.drawImage(ImageReader(img), 50, y_position - new_height, width=new_width, height=new_height)
                    y_position -= new_height + 20
                except Exception as img_error:
                    print(f"无法添加图片 {img_name} 到PDF: {str(img_error)}")
                    c.drawString(50, y_position, "图片无法显示")
                    y_position -= 30

                # 处理描述内容
                if description:
                    paragraphs = description.split('\n')
                    for paragraph in paragraphs:
                        if not paragraph.strip():
                            continue

                        para = Paragraph(paragraph, styles['default'])
                        w, h = para.wrap(width - 100, height)

                        if y_position - h < 50:
                            c.showPage()
                            y_position = height - 40

                        para.drawOn(c, 50, y_position - h)
                        y_position -= h + 10
                else:
                    c.drawString(50, y_position, "未能获取图片描述")
            except Exception as e:
                print(f"处理图片 {os.path.basename(img_path)} 时出错: {str(e)}")

        c.save()
        print(f"PDF报告已生成: {output_pdf}")
    except Exception as e:
        print(f"创建PDF失败: {str(e)}")


def main():
    print("图片场景分析程序开始运行...")

    # 检查API密钥
    if not API_KEY:
        print("警告: 未设置API密钥。请设置环境变量 'doubao_API_KEY'")

    # 设置中文字体
    chinese_font = setup_chinese_font()

    # 获取图片文件列表
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.webp']
    image_files = [os.path.join(IMAGE_DIRECTORY, f) for f in os.listdir(IMAGE_DIRECTORY)
                   if any(f.lower().endswith(ext) for ext in image_extensions)]

    if not image_files:
        print(f"在目录 {IMAGE_DIRECTORY} 中未找到图片文件")
        return

    print(f"找到 {len(image_files)} 个图片文件，开始处理...")

    # 处理每个图片
    image_descriptions = []
    for i, img_path in enumerate(image_files):
        print(f"处理图片 {i + 1}/{len(image_files)}: {os.path.basename(img_path)}")
        description = analyze_image_with_doubao(img_path)
        image_descriptions.append((img_path, description))

    # 创建PDF
    if image_descriptions:
        print("创建PDF报告...")
        create_pdf_with_images(image_descriptions, OUTPUT_PDF, chinese_font)
        print(f"PDF报告已保存至: {OUTPUT_PDF}")
    else:
        print("没有有效的图片描述，不创建PDF")

    print("程序执行完毕")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    except Exception as e:
        print(f"程序执行过程中出现错误: {str(e)}")